require 'rspec'
require 'msf/core/exploit/remote/smb/relay/target_list'

RSpec.describe Msf::Exploit::Remote::SMB::Relay::TargetList do
  let(:subject) { described_class.new '192.0.2.1 192.0.2.2 192.0.2.3', randomize_targets: false }
  let(:user_one) { 'domain/one' }
  let(:user_two) { 'domain/two' }

  describe '#next' do
    context 'when no targets have successfully been relayed to' do
      it 'indefinitely cycles through available targets when there is no identifier provided' do
        3.times do
          expect(subject.next(nil).to_h).to include({ ip: '192.0.2.1', port: 445, protocol: :smb })
          expect(subject.next(nil).to_h).to include({ ip: '192.0.2.2', port: 445, protocol: :smb })
          expect(subject.next(nil).to_h).to include({ ip: '192.0.2.3', port: 445, protocol: :smb })
        end
      end

      it 'cycles through available targets when there is an identifier provided until all relays are in progress' do
        expect(subject.next(user_one).to_h).to include({ ip: '192.0.2.1', port: 445, protocol: :smb })
        expect(subject.next(user_one).to_h).to include({ ip: '192.0.2.2', port: 445, protocol: :smb })
        expect(subject.next(user_one).to_h).to include({ ip: '192.0.2.3', port: 445, protocol: :smb })
        # All relay attempts are in progress, don't return a host
        expect(subject.next(user_one)).to be_nil
      end
    end

    context 'when targets have successfully been relayed to' do
      before(:each) do
        first_target = subject.next(user_one)
        expect(first_target.to_h).to include({ ip: '192.0.2.1', port: 445, protocol: :smb })
        subject.on_relay_end(first_target, identity: user_one, is_success: true)
      end

      it 'no longer returns resolved hosts for a previously successfully relayed identity' do
        expect(subject.next(user_one).to_h).to include({ ip: '192.0.2.2', port: 445, protocol: :smb })
        expect(subject.next(user_one).to_h).to include({ ip: '192.0.2.3', port: 445, protocol: :smb })
        # All relay attempts are inflight, don't return a host
        expect(subject.next(user_one)).to be_nil
      end

      it 'returns the next host to target with a nil identity' do
        3.times do
          expect(subject.next(nil).to_h).to include({ ip: '192.0.2.1', port: 445, protocol: :smb })
          expect(subject.next(nil).to_h).to include({ ip: '192.0.2.2', port: 445, protocol: :smb })
          expect(subject.next(nil).to_h).to include({ ip: '192.0.2.3', port: 445, protocol: :smb })
        end
      end

      it 'returns hosts that have not yet been relayed to for a new identifier' do
        expect(subject.next(user_two).to_h).to include({ ip: '192.0.2.1', port: 445, protocol: :smb })
        expect(subject.next(user_two).to_h).to include({ ip: '192.0.2.2', port: 445, protocol: :smb })
        expect(subject.next(user_two).to_h).to include({ ip: '192.0.2.3', port: 445, protocol: :smb })
        # All relay attempts are inflight, don't return a host
        expect(subject.next(user_two)).to be_nil
      end

      it 'no returns targets were if they were unsuccessfully relayed to' do
        3.times do
          second_target = subject.next(user_one)
          expect(second_target.to_h).to include({ ip: '192.0.2.2', port: 445, protocol: :smb })
          subject.on_relay_end(second_target, identity: user_one, is_success: false)

          third_target = subject.next(user_one)
          expect(third_target.to_h).to include({ ip: '192.0.2.3', port: 445, protocol: :smb })
          subject.on_relay_end(third_target, identity: user_one, is_success: false)
        end
      end

      it 'no longer returns a target if all targets were successfully relayed to that identity' do
        second_target = subject.next(user_one)
        expect(second_target.to_h).to include({ ip: '192.0.2.2', port: 445, protocol: :smb })
        subject.on_relay_end(second_target, identity: user_one, is_success: true)

        third_target = subject.next(user_one)
        expect(third_target.to_h).to include({ ip: '192.0.2.3', port: 445, protocol: :smb })
        subject.on_relay_end(third_target, identity: user_one, is_success: true)

        expect(subject.next(user_one)).to be_nil
      end
    end

    context 'when a target fails being relayed' do
      before(:each) do
        first_target = subject.next(user_one)
        expect(first_target.to_h).to include({ ip: '192.0.2.1', port: 445, protocol: :smb })
        subject.on_relay_end(first_target, identity: user_one, is_success: false)
      end

      it 'can return the same host for the same identity again in the future' do
        expect(subject.next(user_one).to_h).to include({ ip: '192.0.2.2', port: 445, protocol: :smb })
        expect(subject.next(user_one).to_h).to include({ ip: '192.0.2.3', port: 445, protocol: :smb })
        # The host can be tried again
        expect(subject.next(user_one).to_h).to include({ ip: '192.0.2.1', port: 445, protocol: :smb })
        # But the first two hosts are now inflight, and can't be relayed to yet - nil is returned
        expect(subject.next(user_one)).to be_nil
      end
    end
  end
end
